"""
数据源定义模块

定义各种数据源的基类和具体实现。
"""

from abc import ABC, abstractmethod
from enum import Enum
from typing import Dict, List, Any, Optional, Union
from datetime import datetime, date
import pandas as pd
from dataclasses import dataclass


class DataSourceType(Enum):
    """数据源类型"""
    MARKET = "market"           # 市场数据
    FUNDAMENTAL = "fundamental" # 基本面数据
    ALTERNATIVE = "alternative" # 另类数据
    NEWS = "news"              # 新闻数据
    SOCIAL = "social"          # 社交媒体数据
    ECONOMIC = "economic"      # 经济数据


class DataFrequency(Enum):
    """数据频率"""
    TICK = "tick"              # 逐笔数据
    MINUTE = "1min"            # 分钟数据
    MINUTE_5 = "5min"          # 5分钟数据
    MINUTE_15 = "15min"        # 15分钟数据
    MINUTE_30 = "30min"        # 30分钟数据
    HOUR = "1h"                # 小时数据
    DAILY = "1d"               # 日数据
    WEEKLY = "1w"              # 周数据
    MONTHLY = "1m"             # 月数据
    QUARTERLY = "1q"           # 季度数据
    YEARLY = "1y"              # 年数据


@dataclass
class DataRequest:
    """数据请求"""
    symbols: List[str]
    start_date: Union[str, date, datetime]
    end_date: Union[str, date, datetime]
    frequency: DataFrequency
    fields: Optional[List[str]] = None
    filters: Optional[Dict[str, Any]] = None


@dataclass
class DataResponse:
    """数据响应"""
    data: pd.DataFrame
    metadata: Dict[str, Any]
    source: str
    timestamp: datetime
    success: bool
    error_message: Optional[str] = None


class DataSource(ABC):
    """数据源基类"""
    
    def __init__(self, name: str, source_type: DataSourceType, config: Optional[Dict[str, Any]] = None):
        self.name = name
        self.source_type = source_type
        self.config = config or {}
        self.is_connected = False
        self.last_update = None
        self.error_count = 0
        self.success_count = 0
    
    @abstractmethod
    async def connect(self) -> bool:
        """连接数据源"""
        pass
    
    @abstractmethod
    async def disconnect(self) -> bool:
        """断开数据源连接"""
        pass
    
    @abstractmethod
    async def fetch_data(self, request: DataRequest) -> DataResponse:
        """获取数据"""
        pass
    
    @abstractmethod
    async def get_available_symbols(self) -> List[str]:
        """获取可用的股票代码列表"""
        pass
    
    @abstractmethod
    async def get_available_fields(self) -> List[str]:
        """获取可用的字段列表"""
        pass
    
    async def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        try:
            # 简单的连接测试
            if not self.is_connected:
                await self.connect()
            
            return {
                "status": "healthy" if self.is_connected else "unhealthy",
                "name": self.name,
                "type": self.source_type.value,
                "last_update": self.last_update,
                "error_count": self.error_count,
                "success_count": self.success_count,
                "success_rate": self.success_count / max(1, self.success_count + self.error_count)
            }
        except Exception as e:
            return {
                "status": "error",
                "name": self.name,
                "type": self.source_type.value,
                "error": str(e)
            }
    
    def _record_success(self):
        """记录成功"""
        self.success_count += 1
        self.last_update = datetime.now()
    
    def _record_error(self):
        """记录错误"""
        self.error_count += 1


class MarketDataSource(DataSource):
    """市场数据源"""
    
    def __init__(self, name: str, config: Optional[Dict[str, Any]] = None):
        super().__init__(name, DataSourceType.MARKET, config)
        self.supported_frequencies = [
            DataFrequency.TICK,
            DataFrequency.MINUTE,
            DataFrequency.MINUTE_5,
            DataFrequency.MINUTE_15,
            DataFrequency.MINUTE_30,
            DataFrequency.HOUR,
            DataFrequency.DAILY
        ]
        self.supported_fields = [
            "open", "high", "low", "close", "volume", 
            "amount", "vwap", "adj_close"
        ]
    
    async def connect(self) -> bool:
        """连接市场数据源"""
        try:
            # 模拟连接逻辑
            self.is_connected = True
            return True
        except Exception as e:
            self.is_connected = False
            self._record_error()
            return False
    
    async def disconnect(self) -> bool:
        """断开连接"""
        self.is_connected = False
        return True
    
    async def fetch_data(self, request: DataRequest) -> DataResponse:
        """获取市场数据"""
        try:
            if not self.is_connected:
                await self.connect()
            
            # 模拟数据获取
            data = self._generate_mock_market_data(request)
            
            self._record_success()
            return DataResponse(
                data=data,
                metadata={
                    "symbols": request.symbols,
                    "frequency": request.frequency.value,
                    "fields": request.fields or self.supported_fields,
                    "source": self.name
                },
                source=self.name,
                timestamp=datetime.now(),
                success=True
            )
        except Exception as e:
            self._record_error()
            return DataResponse(
                data=pd.DataFrame(),
                metadata={},
                source=self.name,
                timestamp=datetime.now(),
                success=False,
                error_message=str(e)
            )
    
    async def get_available_symbols(self) -> List[str]:
        """获取可用股票代码"""
        return [
            "000001.SZ", "000002.SZ", "000858.SZ", "002415.SZ",
            "600000.SH", "600036.SH", "600519.SH", "600887.SH",
            "AAPL", "GOOGL", "MSFT", "TSLA"
        ]
    
    async def get_available_fields(self) -> List[str]:
        """获取可用字段"""
        return self.supported_fields
    
    def _generate_mock_market_data(self, request: DataRequest) -> pd.DataFrame:
        """生成模拟市场数据"""
        import numpy as np
        
        # 生成时间序列
        dates = pd.date_range(
            start=request.start_date,
            end=request.end_date,
            freq='D'
        )
        
        data_list = []
        for symbol in request.symbols:
            for date in dates:
                # 生成模拟价格数据
                base_price = 100 + hash(symbol) % 100
                open_price = base_price + np.random.normal(0, 2)
                high_price = open_price + abs(np.random.normal(0, 3))
                low_price = open_price - abs(np.random.normal(0, 3))
                close_price = open_price + np.random.normal(0, 2)
                volume = int(np.random.lognormal(10, 1))
                
                data_list.append({
                    'symbol': symbol,
                    'date': date,
                    'open': round(open_price, 2),
                    'high': round(high_price, 2),
                    'low': round(low_price, 2),
                    'close': round(close_price, 2),
                    'volume': volume,
                    'amount': round(close_price * volume, 2)
                })
        
        return pd.DataFrame(data_list)


class FundamentalDataSource(DataSource):
    """基本面数据源"""
    
    def __init__(self, name: str, config: Optional[Dict[str, Any]] = None):
        super().__init__(name, DataSourceType.FUNDAMENTAL, config)
        self.supported_fields = [
            "revenue", "net_income", "total_assets", "total_equity",
            "pe_ratio", "pb_ratio", "roe", "roa", "debt_ratio"
        ]
    
    async def connect(self) -> bool:
        """连接基本面数据源"""
        try:
            self.is_connected = True
            return True
        except Exception:
            self.is_connected = False
            self._record_error()
            return False
    
    async def disconnect(self) -> bool:
        """断开连接"""
        self.is_connected = False
        return True
    
    async def fetch_data(self, request: DataRequest) -> DataResponse:
        """获取基本面数据"""
        try:
            if not self.is_connected:
                await self.connect()
            
            # 模拟数据获取
            data = self._generate_mock_fundamental_data(request)
            
            self._record_success()
            return DataResponse(
                data=data,
                metadata={
                    "symbols": request.symbols,
                    "fields": request.fields or self.supported_fields,
                    "source": self.name
                },
                source=self.name,
                timestamp=datetime.now(),
                success=True
            )
        except Exception as e:
            self._record_error()
            return DataResponse(
                data=pd.DataFrame(),
                metadata={},
                source=self.name,
                timestamp=datetime.now(),
                success=False,
                error_message=str(e)
            )
    
    async def get_available_symbols(self) -> List[str]:
        """获取可用股票代码"""
        return [
            "000001.SZ", "000002.SZ", "000858.SZ", "002415.SZ",
            "600000.SH", "600036.SH", "600519.SH", "600887.SH"
        ]
    
    async def get_available_fields(self) -> List[str]:
        """获取可用字段"""
        return self.supported_fields
    
    def _generate_mock_fundamental_data(self, request: DataRequest) -> pd.DataFrame:
        """生成模拟基本面数据"""
        import numpy as np
        
        data_list = []
        for symbol in request.symbols:
            # 生成模拟财务数据
            revenue = np.random.lognormal(15, 1)  # 营收
            net_income = revenue * np.random.uniform(0.05, 0.25)  # 净利润
            total_assets = revenue * np.random.uniform(2, 5)  # 总资产
            total_equity = total_assets * np.random.uniform(0.3, 0.7)  # 总权益
            
            data_list.append({
                'symbol': symbol,
                'date': datetime.now().date(),
                'revenue': round(revenue, 2),
                'net_income': round(net_income, 2),
                'total_assets': round(total_assets, 2),
                'total_equity': round(total_equity, 2),
                'pe_ratio': round(np.random.uniform(10, 50), 2),
                'pb_ratio': round(np.random.uniform(1, 10), 2),
                'roe': round(net_income / total_equity * 100, 2),
                'roa': round(net_income / total_assets * 100, 2),
                'debt_ratio': round((total_assets - total_equity) / total_assets * 100, 2)
            })
        
        return pd.DataFrame(data_list)


class AlternativeDataSource(DataSource):
    """另类数据源"""
    
    def __init__(self, name: str, config: Optional[Dict[str, Any]] = None):
        super().__init__(name, DataSourceType.ALTERNATIVE, config)
        self.supported_fields = [
            "sentiment_score", "news_count", "social_mentions",
            "search_volume", "analyst_rating", "insider_trading"
        ]
    
    async def connect(self) -> bool:
        """连接另类数据源"""
        try:
            self.is_connected = True
            return True
        except Exception:
            self.is_connected = False
            self._record_error()
            return False
    
    async def disconnect(self) -> bool:
        """断开连接"""
        self.is_connected = False
        return True
    
    async def fetch_data(self, request: DataRequest) -> DataResponse:
        """获取另类数据"""
        try:
            if not self.is_connected:
                await self.connect()
            
            # 模拟数据获取
            data = self._generate_mock_alternative_data(request)
            
            self._record_success()
            return DataResponse(
                data=data,
                metadata={
                    "symbols": request.symbols,
                    "fields": request.fields or self.supported_fields,
                    "source": self.name
                },
                source=self.name,
                timestamp=datetime.now(),
                success=True
            )
        except Exception as e:
            self._record_error()
            return DataResponse(
                data=pd.DataFrame(),
                metadata={},
                source=self.name,
                timestamp=datetime.now(),
                success=False,
                error_message=str(e)
            )
    
    async def get_available_symbols(self) -> List[str]:
        """获取可用股票代码"""
        return [
            "000001.SZ", "000002.SZ", "000858.SZ", "002415.SZ",
            "600000.SH", "600036.SH", "600519.SH", "600887.SH",
            "AAPL", "GOOGL", "MSFT", "TSLA"
        ]
    
    async def get_available_fields(self) -> List[str]:
        """获取可用字段"""
        return self.supported_fields
    
    def _generate_mock_alternative_data(self, request: DataRequest) -> pd.DataFrame:
        """生成模拟另类数据"""
        import numpy as np
        
        data_list = []
        for symbol in request.symbols:
            data_list.append({
                'symbol': symbol,
                'date': datetime.now().date(),
                'sentiment_score': round(np.random.uniform(-1, 1), 3),
                'news_count': int(np.random.poisson(10)),
                'social_mentions': int(np.random.poisson(50)),
                'search_volume': int(np.random.lognormal(8, 1)),
                'analyst_rating': round(np.random.uniform(1, 5), 1),
                'insider_trading': round(np.random.normal(0, 1000000), 2)
            })
        
        return pd.DataFrame(data_list)


# 预定义的数据源实例
DEFAULT_MARKET_SOURCE = MarketDataSource("默认市场数据源")
DEFAULT_FUNDAMENTAL_SOURCE = FundamentalDataSource("默认基本面数据源")
DEFAULT_ALTERNATIVE_SOURCE = AlternativeDataSource("默认另类数据源")

# 数据源注册表
DATA_SOURCE_REGISTRY = {
    "market_default": DEFAULT_MARKET_SOURCE,
    "fundamental_default": DEFAULT_FUNDAMENTAL_SOURCE,
    "alternative_default": DEFAULT_ALTERNATIVE_SOURCE
}


def get_data_source(name: str) -> Optional[DataSource]:
    """获取数据源"""
    return DATA_SOURCE_REGISTRY.get(name)


def register_data_source(name: str, source: DataSource):
    """注册数据源"""
    DATA_SOURCE_REGISTRY[name] = source


def list_data_sources() -> Dict[str, DataSource]:
    """列出所有数据源"""
    return DATA_SOURCE_REGISTRY.copy()